#light

type Distribution<'a> =
    abstract Sample : 'a
    abstract Support : Set<'a>
    abstract Expectation : ('a -> float) -> float
    
let always x =
    { new Distribution<'a> with
        member d.Sample = x
        member d.Support = Set.singleton x
        member d.Expectation(H) = H(x) }

let rnd = System.Random()

let coinFlip (p:float) (d1:Distribution<'a>) (d2:Distribution<'a>) =
    if p < 0.0 || p > 1.0 then failwith "Invalid probability in coinFlip"
    { new Distribution<'a> with
        member d.Sample =
            if rnd.NextDouble() < p then d1.Sample else d2.Sample
        member d.Support = Set.Union(d1.Support,d2.Support)
        member d.Expectation(H) =
            p * d1.Expectation(H) + (1.0-p) * d2.Expectation(H) }
        
let bind (dist:Distribution<'a>) (k: 'a -> Distribution<'b>) =
    { new Distribution<'b> with
        member d.Sample = (k(dist.Sample)).Sample
        member d.Support = Set.Union(dist.Support.Map(fun d -> (k d).Support))
        member d.Expectation(H) = dist.Expectation(fun x -> (k x).Expectation(H)) }
        
type DistributionBuilder() =
    member x.Delay(f) = bind (always ()) f
    member x.Let(v,f) = bind (always v) f
    member x.Bind(d,f) = bind d f
    member x.Return(x) = always x
let dist = new DistributionBuilder()
    
    
let weightedCases (inp: ('a * float) list) =
    let rec coinFlips w l =
        match l with
        | []          -> failwith "no coinFlips"
        | [(d,_)]     -> always d
        | (d,p)::rest -> coinFlip (p/(1.0-w)) (always d) (coinFlips (w+p) rest)
    coinFlips 0.0 inp

let countedCases inp =
    let total = List.sumByInt (fun (_,v) -> v) inp
    weightedCases (inp.Map (fun (x,v) -> (x,(float v/float total))))


// Calculate roulette odds
type Outcome = Even | Odd | Zero
let roulette = countedCases [ Even,18; Odd,18; Zero,1]

// draw from this distribution
roulette.Sample

roulette.Sample

// compute the expected payout of a $5 bet on Even, where we would get a $10 return
roulette.Expectation (function Even -> 10.0 | Odd -> 0.0 | Zero -> 0.0)





